import re
import json
from typing import List, Dict, Union

class LegalTextSplitter:
    """法律条文文本拆分器"""
    
    def __init__(self):
        # 匹配法律条文的模式：第X条 或 第X条 开头的内容
        self.article_pattern = re.compile(r'^第[一二三四五六七八九十百千万\d]+条\s*.*', re.MULTILINE)
        
        # 匹配章节标题的模式：第X章 或 第X节
        self.chapter_pattern = re.compile(r'^第[一二三四五六七八九十百千万\d]+[章节]\s*.*', re.MULTILINE)
        
        # 匹配序言、总纲等特殊章节
        self.special_sections = ['序言', '总纲', '公民的基本权利和义务', '国家机构']
    
    def split_legal_text(self, text: str) -> List[Dict[str, Union[str, str, int]]]:
        """
        拆分法律条文文本
        
        Args:
            text: 法律条文文本内容
            
        Returns:
            拆分后的文本单元列表，每个单元包含type、content、index等字段
        """
        # 预处理文本，统一换行符
        text = text.replace('\r\n', '\n').replace('\r', '\n')
        
        # 按行分割文本
        lines = text.split('\n')
        
        units = []
        current_unit = {"type": "", "content": "", "index": 0}
        unit_index = 0
        
        i = 0
        while i < len(lines):
            line = lines[i].strip()
            
            # 跳过空行
            if not line:
                i += 1
                continue
            
            # 检查是否是章节标题
            if self.chapter_pattern.match(line):
                # 保存之前的单元
                if current_unit["content"]:
                    current_unit["index"] = unit_index
                    units.append(current_unit.copy())
                    unit_index += 1
                
                # 开始新的章节单元
                current_unit = {
                    "type": "chapter",
                    "content": line,
                    "index": unit_index
                }
                units.append(current_unit.copy())
                unit_index += 1
                current_unit = {"type": "", "content": "", "index": 0}
                
            # 检查是否是法律条文
            elif self.article_pattern.match(line):
                # 保存之前的单元
                if current_unit["content"]:
                    current_unit["index"] = unit_index
                    units.append(current_unit.copy())
                    unit_index += 1
                
                # 开始新的条文单元
                article_content = line
                
                # 查找条文的完整内容（可能跨多行）
                j = i + 1
                while j < len(lines):
                    next_line = lines[j].strip()
                    # 如果下一行是空行或者是新的条文/章节，则停止
                    if not next_line or self.article_pattern.match(next_line) or self.chapter_pattern.match(next_line):
                        break
                    # 如果下一行以数字或中文数字开头，可能是条文的子项
                    if re.match(r'^[（\(]?[一二三四五六七八九十\d]+[）\)]?', next_line):
                        article_content += '\n' + next_line
                    # 如果下一行是条文的继续内容
                    elif next_line and not next_line.startswith('第'):
                        article_content += '\n' + next_line
                    j += 1
                
                current_unit = {
                    "type": "article",
                    "content": article_content,
                    "index": unit_index
                }
                units.append(current_unit.copy())
                unit_index += 1
                current_unit = {"type": "", "content": "", "index": 0}
                i = j - 1  # 调整索引
                
            # 检查是否是特殊章节（序言、总纲等）
            elif any(section in line for section in self.special_sections):
                # 保存之前的单元
                if current_unit["content"]:
                    current_unit["index"] = unit_index
                    units.append(current_unit.copy())
                    unit_index += 1
                
                # 开始新的特殊章节单元
                section_content = line
                
                # 查找章节的完整内容
                j = i + 1
                while j < len(lines):
                    next_line = lines[j].strip()
                    # 如果遇到新的章节或条文，则停止
                    if self.chapter_pattern.match(next_line) or self.article_pattern.match(next_line):
                        break
                    if next_line:
                        section_content += '\n' + next_line
                    j += 1
                
                current_unit = {
                    "type": "section",
                    "content": section_content,
                    "index": unit_index
                }
                units.append(current_unit.copy())
                unit_index += 1
                current_unit = {"type": "", "content": "", "index": 0}
                i = j - 1  # 调整索引
                
            # 普通段落内容
            else:
                if current_unit["type"] == "":
                    current_unit["type"] = "paragraph"
                    current_unit["content"] = line
                else:
                    current_unit["content"] += '\n' + line
            
            i += 1
        
        # 保存最后一个单元
        if current_unit["content"]:
            current_unit["index"] = unit_index
            units.append(current_unit.copy())
        
        return units
    
    def save_to_json(self, units: List[Dict], output_file: str):
        """保存拆分结果到JSON文件"""
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(units, f, ensure_ascii=False, indent=2)
    
    def save_to_txt(self, units: List[Dict], output_file: str):
        """保存拆分结果到TXT文件"""
        with open(output_file, 'w', encoding='utf-8') as f:
            for unit in units:
                f.write(f"=== {unit['type'].upper()} {unit['index']} ===\n")
                f.write(unit['content'])
                f.write('\n\n')
    
    def print_summary(self, units: List[Dict]):
        """打印拆分结果摘要"""
        print(f"总共拆分为 {len(units)} 个单元")
        
        type_counts = {}
        for unit in units:
            unit_type = unit['type']
            type_counts[unit_type] = type_counts.get(unit_type, 0) + 1
        
        print("\n各类型单元统计：")
        for unit_type, count in type_counts.items():
            print(f"  {unit_type}: {count} 个")


def main():
    """主函数示例"""
    # 创建拆分器
    splitter = LegalTextSplitter()
    
    # 读取宪法文件
    try:
        with open('宪法.md', 'r', encoding='utf-8') as f:
            text = f.read()
    except FileNotFoundError:
        print("错误：找不到 宪法.md 文件")
        return
    
    # 拆分文本
    print("正在拆分法律条文...")
    units = splitter.split_legal_text(text)
    
    # 打印摘要
    splitter.print_summary(units)
    
    # 保存结果
    splitter.save_to_json(units, '宪法_拆分结果.json')
    splitter.save_to_txt(units, '宪法_拆分结果.txt')
    
    print("\n拆分完成！")
    print("结果已保存到：")
    print("  - 宪法_拆分结果.json (JSON格式)")
    print("  - 宪法_拆分结果.txt (文本格式)")
    
    # 显示前几个单元作为示例
    print("\n前5个单元示例：")
    for i, unit in enumerate(units[:5]):
        print(f"\n--- 单元 {i+1} ---")
        print(f"类型: {unit['type']}")
        print(f"内容: {unit['content'][:100]}...")


if __name__ == "__main__":
    main() 